# condition_checker.py
import pandas as pd

from pick.condition import Condition


class Strategy:
    def check(self, data):
        raise NotImplementedError("This method should be overridden by subclasses")


class PriceAboveStrategy(Strategy):
    """
      简单的价格模式，具体的筛选策略，自己实现，这是一个简单的示例
     data的数据格式 ["timestamp", "open", "high", "low", "close", "confirm"]
    """
    def __init__(self, condition: Condition):
        self.condition = condition

    # 判断最近连续3根K线是否是光头阳线
    def _is_three_above_ema20(self, df):
        last_3 = df.head(3)  # 获取最近3根K线
        # 如果需要比较最近3个bar和EMA20
        if self.condition.recent_three_bars_compare_ema20:
            if self.condition.direction == 1:
                # 如果是上涨方向，所有收盘价要高于EMA20
                return (last_3['close'] >= last_3['EMA20']).all()

            else:
                # 如果是下跌方向，所有收盘价要低于EMA20
                return (last_3['close'] <= last_3['EMA20']).all()

        return True
        # return ((last_3['close'] > last_3['open']) &  # 阳线：收盘价 > 开盘价
        #         (last_3['high'] == last_3['close']) &  # 没有上影线：最高价 == 收盘价
        #         (last_3['low'] == last_3['open'])).all()  # 没有下影线：最低价 == 开盘价

    def _is_three_consecutive(self, df):
        last_3 = df.head(3)  # 获取最近3根K线
        # 如果需要比较最近3个bar具有连续性
        if self.condition.three_bars_consecutive:
            if self.condition.direction == 1:
                # 如果是上涨方向，所有收盘价要高于EMA20
                return (last_3['close'] >= last_3['open']).all()
            else:
                return (last_3['close'] <= last_3['open']).all()
        return True



    # 判断最近的K线是否都在EMA20上方
    def _is_all_above_ema20(self, df):
        last_n = df.head(self.condition.consecutive_bars)
        # 如果需要比较最近3个bar和EMA20

        if self.condition.direction == 1:
            return (last_n['close'] >= last_n['EMA20']).all()  # 判断是否每根K线的收盘价高于EMA20
        else:
            return (last_n['close'] <= last_n['EMA20']).all()  # 判断是否每根K线的收盘价高于EMA20

    import pandas as pd

    def analyze_kline_dataframe(self, df):
        """
        分析K线特性
        参数:
            df (DataFrame): 包含K线数据，列包括 'open', 'close', 'high', 'low'
        返回:
            bool: 是否符合特性要求
        """
        # 取前6条数据并按时间升序排序
        selected_df = df.head(6).iloc[::-1].reset_index(drop=True)

        # 获取第2根K线的开盘价和收盘价
        second_open = selected_df.iloc[1]['open']
        second_close = selected_df.iloc[1]['close']

        # 前两根K线验证是否是大阳线
        for i in range(2):
            if selected_df.iloc[i]['close'] <= selected_df.iloc[i]['open']:
                return False  # 第i+1根K线不是大阳线
            close_price = selected_df.iloc[i]['close']
            open_price = selected_df.iloc[i]['open']
            # 检查振幅是否大于2%
            amplitude = (close_price - open_price) / open_price * 100
            if amplitude <= 2:
                return False  # 振幅规则不满足

        # 检查后续4根K线（第3至第6根）
        for i in range(2, 6):
            open_price = selected_df.iloc[i]['open']
            close_price = selected_df.iloc[i]['close']
            high_price = selected_df.iloc[i]['high']
            low_price = selected_df.iloc[i]['low']

            # 检查阴阳规则
            if i % 2 == 0:  # 偶数序号：阴线
                if close_price >= open_price:
                    return False  # 阴线规则不满足
            else:  # 奇数序号：阳线
                if close_price <= open_price:
                    return False  # 阳线规则不满足

            # 检查振幅是否大于2%
            amplitude = (high_price - low_price) / low_price * 100
            if amplitude <= 2:
                return False  # 振幅规则不满足

        return True  # 所有条件满足

    def check(self, data):
        return self.analyze_kline_dataframe(data)
        # return self._is_all_above_ema20(data) and self._is_three_above_ema20(data) and self._is_three_consecutive(data)




# 构造一些股票 k线数据 DataFrame
df = pd.DataFrame([
    ['2023-07-01 18:30:00', 0.8847, 0.9280, 0.8765, 0.9223, 1],
    ['2023-07-01 17:35:00', 0.9074, 0.9314, 0.8667, 0.8851, 1],
    ['2023-07-01 16:40:00', 0.8577, 0.9255, 0.8388, 0.9081, 1],
    ['2023-07-01 15:45:00', 0.9430, 0.9574, 0.8575, 0.8577, 1],
    ['2023-07-01 14:50:00', 0.8589, 0.9623, 0.8552, 0.9431, 1],
    ['2023-07-01 13:55:00', 0.7966, 0.9004, 0.7914, 0.8589, 1],
    ['2023-07-01 12:00:00', 0.7806, 0.8003, 0.7636, 0.7967, 1],
    ['2023-07-01 10:05:00', 0.7832, 0.7882, 0.7576, 0.7804, 1]
])
df.columns = ['timestamp', 'open', 'high', 'low', 'close', 'confirm']


def analyze_kline_dataframe(df):
    """
    分析K线特性
    参数:
        df (DataFrame): 包含K线数据，列包括 'open', 'close', 'high', 'low'
    返回:
        bool: 是否符合特性要求
    """
    # 取前6条数据并按时间升序排序
    selected_df = df.head(6).iloc[::-1].reset_index(drop=True)

    # 获取第2根K线的开盘价和收盘价
    second_open = selected_df.iloc[1]['open']
    second_close = selected_df.iloc[1]['close']

    # 前两根K线验证是否是大阳线
    for i in range(2):
        if selected_df.iloc[i]['close'] <= selected_df.iloc[i]['open']:
            return False  # 第i+1根K线不是大阳线
        close_price = selected_df.iloc[i]['close']
        open_price = selected_df.iloc[i]['open']
        # 检查振幅是否大于2%
        amplitude = (close_price - open_price) / open_price * 100
        if amplitude <= 2:
            return False  # 振幅规则不满足
    # 检查后续4根K线（第3至第6根）
    for i in range(2, 6):
        open_price = selected_df.iloc[i]['open']
        close_price = selected_df.iloc[i]['close']
        high_price = selected_df.iloc[i]['high']
        low_price = selected_df.iloc[i]['low']

        # 检查阴阳规则
        if i % 2 == 0:  # 偶数序号：阴线
            if close_price >= open_price:
                return False  # 阴线规则不满足
        else:  # 奇数序号：阳线
            if close_price <= open_price:
                return False  # 阳线规则不满足

        # 检查振幅是否大于2%
        amplitude = (high_price - low_price) / low_price * 100
        if amplitude <= 2:
            return False  # 振幅规则不满足

    return True  # 所有条件满足



print(analyze_kline_dataframe(df))